﻿module TypeEnvironment
open Ast
open System.Collections.Generic
open SymbolTableTypes
open System
open MaybeMonad


type Maybe<'a> = Environment -> ('a * Environment) MaybeOrFail

and Environment = 
  {
    mutable LineNumber : int
    Tables : List<GenericTable>
    WhileCounter : int ref
    AliasCounter : int ref
  }

type TypeCheckState = Maybe<Ast.Type>

(*
State monad builder.
*)
//let state = new StateMonad() 
let maybe = MaybeBuilder()

let fail error = fun s -> Fail (error + " at line " + (s.LineNumber |> string))
(*
Monadic state getter.
*)
let private getState = (fun s -> Just(s, s)) 

let reset_alias_counter =
  fun s ->
    Just ((), {s with AliasCounter = ref 1} )

let ignore_ (_,s) =
  match s with
  | Fail error -> 
    Fail error 
  | _ ->
    Just((), s)

let private get_container_class_st = 
  fun (s : Environment) ->
    Just ((get_class_table s.Tables.Head), s)

let private get_root = 
  fun (s : Environment) ->
    Just ((get_global_table s.Tables.Head), s)

let private get_body_st = 
  fun (s : Environment) ->
    Just ((get_body_table s.Tables.Head), s)

let private get_inherited_list = 
  fun (s : Environment) ->
    Just ((get_inherited_class_table s.Tables.Head), s)

let private get_env = 
  fun (s : Environment) ->
    Just ((s.Tables.Head), s)

let set_line (line, _) =
  fun (s : Environment) ->
    s.LineNumber <- line
    Just ((), s)

let add_block_env block = 
  fun (s : Environment) ->
    match s.Tables.Head with
    | GenericTable.BodyTable b -> 
      let block_table = 
        {
          Parent = Body b
          Symbols = Dictionary<string, BodySymbol>()
          Block = block
          Children = List<BlockTable>()
        }
      do b.Block <- Some block_table
      do s.Tables.Push(GenericTable.BlockTable block_table)
      Just((),s)
    | GenericTable.BlockTable b -> 
      let block_table = 
        {
          Parent = Block b
          Symbols = Dictionary<string, BodySymbol>()
          Block = block
          Children = List<BlockTable>()
        }
      do b.Children.AddLast(block_table) |> ignore
      do s.Tables.Push(GenericTable.BlockTable block_table)
      Just((),s)
    | _ -> 
      Fail "This doesn't have sense"

let add_method_env (meth : ReturnType * id * Arg list * Block) = 
  fun (s : Environment) ->
    let ret, name, args, block = meth
    match s.Tables.Head with
    | GenericTable.ClassTable c -> 
      let body_table = 
        {
          Name          = name
          Container     = c
          Args          = Dictionary<string, BodySymbol>()
          ReturnType    = ret
          Block         = None
          Aliases       = ResizeArray()
        }
      do c.Methods.[name] <-  { Name = name; Infos = ((args |> List.map (fun (Arg (t, _)) -> t)), ret, ref (Some body_table)) }
      do s.Tables.Push(GenericTable.BodyTable body_table)
      Just((),s)
    | _ -> 
      Fail "This doesn't have sense"

let remove_env = 
  fun (s : Environment) ->
    do s.Tables.Pop () |> ignore
    Just((), s)

let add_field_symbol name ty = 
  fun (s : Environment) ->
    match s.Tables.Head with
//    | GenericTable.ClassTable c ->
//      do c.Symbols.Add(name, { Name = name; Kind = Field ty })
//      Just((),s)
    | GenericTable.BodyTable b ->
      let alias = name + "$" + (string s.AliasCounter.Value)
      do incr s.AliasCounter
      do b.Aliases.Add ({ Name = alias; Type = ty; Alias = alias })
      do b.Args.Add(name, { Name = name; Type = ty; Alias = alias })
      Just((),s)
    | GenericTable.BlockTable b as t ->
      let body_table = get_body_table t
      let alias = name + "$" + (string s.AliasCounter.Value)
      do incr s.AliasCounter
      do body_table.Aliases.Add ({ Name = alias; Type = ty; Alias = alias })
      do b.Symbols.Add(name, { Name = name; Type = ty; Alias = alias })

      Just((),s)
    | _ -> Fail "This doesn't have sense"

let add_class_env name = 
  fun (s : Environment) ->
    match s.Tables.Head with
    | GenericTable.GlobalSymbols g ->
      do s.Tables.Push(GenericTable.ClassTable g.[name])
      Just((),s)
    | _ ->
      Fail "This doesn't have sense"

let enter_while = 
  fun (s : Environment) ->
    do incr s.WhileCounter 
    Just((), s)  

let exit_while = 
  fun (s : Environment) ->
    do decr s.WhileCounter 
    Just((), s)     

let check_is_while =
  fun (s : Environment) ->
    if !s.WhileCounter > 0 then
      Just ((), s) 
    else
      Fail "break and continue are allowed only in while blocks"

let get_inherited_list_with name = 
  maybe{
    let! globals = get_root
//    if globals.ContainsKey(name) then
    return (get_inherited_class_table (ClassTable globals.[name]))
//    else
//      return None
  }

let get_container_class_name =
  maybe{
    let! class_st = get_container_class_st
    return class_st.Name
  }

let get_classes = 
  maybe{
    let! global_st = get_root
    return global_st.Keys
  }

let reduce_class pattern error = 
  maybe{
    let! global_st = get_root
    if global_st.Keys |> Seq.exists (pattern) then
      return ()
    else
      return! fail error
  }

let reduce_inherited pattern = 
  maybe{
    let! inherited_list = get_inherited_list
    return inherited_list |> Seq.exists pattern
  }

let try_find_method_with class_name name = 
  maybe{
    let! inherited_list = get_inherited_list_with class_name
//    match inherited_list with
//    | None -> 
//      return None 
//    | Some inherited_list ->
    match inherited_list |> Seq.tryFind (fun c -> c.Methods.ContainsKey(name)) with
    | None -> return! fail (name + " doesn't exist in " + class_name)
    | Some t -> 
      let (args, ret, _) = t.Methods.[name].Infos
      return (ret, args)      
  }

let try_find_method name = 
  maybe{
    let! c_name = get_container_class_name
    return! try_find_method_with c_name name
  }

let try_find_var_with class_name name = 
  maybe{
    let! inherited_list = get_inherited_list_with class_name
    match inherited_list |> Seq.tryFind (fun c -> c.Fields.ContainsKey(name)) with
    | None -> return! fail ("identifier " + name + " doesnt't exist")
    | Some t -> 
      return t.Fields.[name].Type
  }

//let try_find_var name = 
//  maybe{
//    let! s = getState
//    let rec find_tu_class 
//    s.Tables.Head
//    let! c_name = get_container_class_name
//    do Console.WriteLine(string name)
//    return! try_find_var_with c_name name
//  }

let is_arg loc = 
  maybe{
    let! body_st = get_body_st
    match loc with
    | Location.Id name ->
      return body_st.Args.ContainsKey(name)
    | _ -> 
      return false
  }

let get_metod_return_type = 
  maybe{
    let! body_st = get_body_st
    return body_st.ReturnType
  }

let check_field_name name = 
  maybe{
    let! first_env = get_env
    let rec exists (t : GenericTable) = 
      maybe{
        match t with
        | GenericTable.ClassTable c ->
          let inherited_list = get_inherited_class_table t
          if inherited_list |> Seq.exists (fun c -> c.Fields.ContainsKey(name)) then
            return! fail (name + " already exist in class or iherited")
          else
            return ()
        | GenericTable.BodyTable b ->
          if b.Args.ContainsKey(name) then
            return! fail (name + " already exist in args")
          else
            return! exists (GenericTable.ClassTable b.Container)
        | GenericTable.BlockTable b ->
          if b.Symbols.ContainsKey(name) then
            return! fail (name + " already exist in block")
          else
            return! exists (match b.Parent with | BodyOrBlock.Block b -> GenericTable.BlockTable b | BodyOrBlock.Body b -> GenericTable.BodyTable b )
        | GenericTable.GlobalSymbols g ->          
          return! fail "this make nonsense"
      }
    return! exists first_env
  }

let rec find_var name = 
  maybe{
    let! first_env = get_env
    let rec exists (t : GenericTable) = 
      maybe{
        match t with
        | GenericTable.ClassTable c ->
          let inherited_list = get_inherited_class_table t
          match inherited_list |> Seq.tryFind (fun c -> c.Fields.ContainsKey(name)) with
          | None ->
            return! fail (name + " doesn't exist in class or iherited")
          | Some t ->
            return t.Fields.[name].Type
        | GenericTable.BodyTable b ->
          if b.Args.ContainsKey(name) then
            return b.Args.[name].Type
          else
            return! exists (GenericTable.ClassTable b.Container)
        | GenericTable.BlockTable b ->
          if b.Symbols.ContainsKey(name) then
            return b.Symbols.[name].Type 
          else
            return! exists (match b.Parent with | BodyOrBlock.Block b -> GenericTable.BlockTable b | BodyOrBlock.Body b -> GenericTable.BodyTable b )
        | GenericTable.GlobalSymbols g ->          
          return! fail "this make nonsense"
      }
    return! exists first_env
  }

